#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
蓝奏云提取码文件下载器
依赖: pip install lanzou-api
用法: python3.9 lanzou_down.py
"""

import os
import sys
import re
import urllib.request
from urllib.parse import unquote
from lanzou.api import LanZouCloud

# ---------- 工具函数 ----------
def human_size(n):
    for unit in ['B', 'KB', 'MB', 'GB']:
        if n < 1024:
            return f'{n:.2f}{unit}'
        n /= 1024
    return f'{n:.2f}TB'

def download(url: str, save_path: str):
    """支持断点续传的下载"""
    req = urllib.request.Request(url)
    exist_size = 0
    if os.path.exists(save_path):
        exist_size = os.path.getsize(save_path)
        req.add_header('Range', f'bytes={exist_size}-')

    try:
        with urllib.request.urlopen(req) as resp:
            total = int(resp.headers.get('Content-Length', 0)) + exist_size
            print(f'[INFO] 文件大小: {human_size(total)}')
            mode = 'ab' if exist_size else 'wb'
            with open(save_path, mode) as f:
                done = exist_size
                while True:
                    chunk = resp.read(64 * 1024)
                    if not chunk:
                        break
                    f.write(chunk)
                    done += len(chunk)
                    percent = done / total * 100
                    print(f'\r[DOWN] {percent:6.2f}%  {human_size(done)}/{human_size(total)}', end='')
            print('\n[INFO] 下载完成 ->', save_path)
    except Exception as e:
        print('\n[ERR ] 下载失败:', e)
        sys.exit(1)

# ---------- 主逻辑 ----------
def main():
    lzy = LanZouCloud()
    url = input('请输入蓝奏云分享链接: ').strip()
    pwd = input('请输入提取码（无则直接回车）: ').strip() or None

    # 解析
    print('[INFO] 正在解析...')
    code = lzy.get_durl_by_url(url, pwd)
    if code == LanZouCloud.SUCCESS:
        info = lzy.get_share_info_by_url(url, pwd)
        file_name = info.name
        direct_url = info.durl
        print(f'[INFO] 文件名: {file_name}')
        print(f'[INFO] 直链: {direct_url[:60]}...')
        save_name = re.sub(r'[\\/:*?"<>|]', '_', file_name)   # 合法化文件名
        if os.path.exists(save_name):
            print(f'[WARN] 文件已存在，将续传: {save_name}')
        download(direct_url, save_name)
    else:
        print('[ERR ] 解析失败，请检查链接/提取码')
        sys.exit(1)

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.exit('\nCancelled by user')